﻿using CrazyMVSamLibrary.Base;
using CrazyMVSamLibrary.Model;
using Emgu.CV;
using Microsoft.ML.OnnxRuntime.Tensors;
using Microsoft.ML.OnnxRuntime;
using System.Collections.Generic;
using System.Linq;

namespace CrazyMVSamLibrary.Function
{
    internal class Sam
    {
        InferenceSession mEncoder;
        InferenceSession mDecoder;
        public float mask_threshold = 0.0f;

        public void LoadONNXModel(string encode_model_path, string decode_model_path)
        {
            if (this.mEncoder != null)
                this.mEncoder.Dispose();

            if (this.mDecoder != null)
                this.mDecoder.Dispose();

            this.mEncoder = new InferenceSession(encode_model_path);
            this.mDecoder = new InferenceSession(decode_model_path);
        }

        public float[] Encode(Transforms tranform, Mat image, int orgWid, int orgHei)
        {
            float[] img = tranform.ApplyImage(image, orgWid, orgHei);
            DenseTensor<float> tensor = new DenseTensor<float>(img, new[] { 1, 3, 1024, 1024 });
            List<NamedOnnxValue> inputs = new List<NamedOnnxValue>
            {
                NamedOnnxValue.CreateFromTensor(this.mEncoder.InputNames[0], tensor)
            };

            IDisposableReadOnlyCollection<DisposableNamedOnnxValue> results = this.mEncoder.Run(inputs);
            float[] embedding = results.First().AsTensor<float>().ToArray();
            results.Dispose();
            return embedding;
        }

        public MaskData Decode(List<Promotion> promotions, float[] embedding, int orgWid, int orgHei)
        {
            DenseTensor<float> embedding_tensor = new DenseTensor<float>(embedding, new[] { 1, 256, 64, 64 });
            List<Promotion> bpmos = promotions.FindAll(e => e.mType == PromotionType.Box);
            List<Promotion> pproms = promotions.FindAll(e => e.mType == PromotionType.Point);
            int boxCount = promotions.FindAll(e => e.mType == PromotionType.Box).Count();
            int pointCount = promotions.FindAll(e => e.mType == PromotionType.Point).Count();
            float[] promotion = new float[2 * (boxCount * 2 + pointCount)];
            float[] label = new float[boxCount * 2 + pointCount];
            for (int i = 0; i < boxCount; i++)
            {
                float[] input = bpmos[i].GetInput();
                for (int j = 0; j < input.Count(); j++)
                {
                    promotion[4 * i + j] = input[j];
                }
                float[] la = bpmos[i].GetLable();
                for (int j = 0; j < la.Count(); j++)
                {
                    label[2 * i + j] = la[j];
                }
            }
            for (int i = 0; i < pointCount; i++)
            {
                float[] p = pproms[i].GetInput();
                for (int j = 0; j < p.Count(); j++)
                {
                    promotion[boxCount * 4 + 2 * i + j] = p[j];
                }
                float[] la = pproms[i].GetLable();
                for (int j = 0; j < la.Count(); j++)
                {
                    label[boxCount * 2 + i + j] = la[j];
                }
            }
            DenseTensor<float> point_coords_tensor = new DenseTensor<float>(promotion, new[] { 1, boxCount * 2 + pointCount, 2 });
            DenseTensor<float> point_label_tensor = new DenseTensor<float>(label, new[] { 1, boxCount * 2 + pointCount });
            float[] mask = new float[256 * 256];
            for (int i = 0; i < mask.Count(); i++)
            {
                mask[i] = 0;
            }
            DenseTensor<float> mask_tensor = new DenseTensor<float>(mask, new[] { 1, 1, 256, 256 });
            float[] hasMaskValues = new float[1] { 0 };
            DenseTensor<float> hasMaskValues_tensor = new DenseTensor<float>(hasMaskValues, new[] { 1 });
            float[] orig_im_size_values = { (float)orgHei, (float)orgWid };
            DenseTensor<float> orig_im_size_values_tensor = new DenseTensor<float>(orig_im_size_values, new[] { 2 });
            List<NamedOnnxValue> decode_inputs = new List<NamedOnnxValue>
            {
                NamedOnnxValue.CreateFromTensor("image_embeddings", embedding_tensor),
                NamedOnnxValue.CreateFromTensor("point_coords", point_coords_tensor),
                NamedOnnxValue.CreateFromTensor("point_labels", point_label_tensor),
                NamedOnnxValue.CreateFromTensor("mask_input", mask_tensor),
                NamedOnnxValue.CreateFromTensor("has_mask_input", hasMaskValues_tensor),
                NamedOnnxValue.CreateFromTensor("orig_im_size", orig_im_size_values_tensor)
            };
            MaskData md = new MaskData();
            IDisposableReadOnlyCollection<DisposableNamedOnnxValue> segmaskValue = this.mDecoder.Run(decode_inputs);
            List<DisposableNamedOnnxValue> segmask = segmaskValue.ToList();
            md.mMask = segmask[0].AsTensor<float>().ToArray().ToList();
            md.mShape = segmask[0].AsTensor<float>().Dimensions.ToArray();
            md.mIoU = segmask[1].AsTensor<float>().ToList();
            segmaskValue.Dispose();
            return md;
        }

        public void Clear()
        {
            this.mEncoder.Dispose();
            this.mDecoder.Dispose();
        }
    }
}